#include "Optimization.h"
#include "Exception.h"
#include "ScopeDictionary.h"

namespace Compiler
{
	namespace Intermediate
	{
		// This optimizer perform following transformations:
		// 1. Remove instructions that involve identities, e.g. replace
		//	  a = b + 1 with a = b
		// 2. Remove useless phi functions, e.g. a[1] = phi(a[0], a[0])
		// 3. Fold constants, e.g. a = 1 + 2 => a = 3
		class UselessInstructionOptimizer : public IntraProcOptimizer
		{
		private:
			String GetExprString(Variable * var)
			{
				StringBuilder sb(48);
				sb.Append(L'#');
				sb.Append(String(var->Id, 16));
				return sb.ProduceString();
			}
			void GetOperandExprString(StringBuilder & sb, Operand & op)
			{
				if (op.IsIntegral() || op.Type == OperandType::ConstFloat)
					sb.Append(String(op.IntValue, 16));
				else if (op.Type == OperandType::ConstDouble)
				{
					sb.Append(String(op.Values[0], 16));
					sb.Append(String(op.Values[1], 16));
				}
				else if (op.IsVariable())
				{
					sb.Append(L'#');
					sb.Append(String(op.Var->Id, 16));
				}
				else
					throw InvalidProgramException(L"Unknown operand type.");
			}
			String GetExprString(Instruction & instr)
			{
				StringBuilder sb(128);
				if (instr.Func)
				{
					sb.Append(instr.Func->Name);
					sb.Append(L':');
				}
				int i = 0;
				for (auto& op : instr.Operands)
				{
					if (i>0)
						sb.Append(L',');
					GetOperandExprString(sb, op);
				}
				return sb.ProduceString();
			}
			bool TraverseNode(ControlFlowGraph *program, ControlFlowNode * node, ScopeDictionary<String, Operand> & exprMap)
			{
				bool changed = false;
				if (node == program->Sink)
					return false;
				exprMap.PushScope();
				if (node != program->Source)
				{
					for (auto iter = node->Code.begin(); iter!=node->Code.end(); iter++)
					{
						Instruction & instr = iter.Current->Value;
						// Process Phi functions
						if (instr.Func == Operation::Phi)
						{
							Variable * v = instr.Operands[0].Var;
							bool sameArg = true;
							for (auto op : instr.Operands)
							{
								if (op.Var != v && op.Var != instr.LeftOperand.Var)
								{
									sameArg = false;
									break;
								}
							}
							if (sameArg)
							{
								Operand valueOp;
								/*if (!exprMap.TryGetValue(GetExprString(v), valueOp))
									throw InvalidProgramException(L"Value entry not found in exprMap");*/
								exprMap.Add(GetExprString(instr.LeftOperand.Var), instr.Operands[0]);
								iter.Current->Delete();
								changed = true;
								continue;
							}
							else
							{
								String exprString = GetExprString(instr);
								Operand opValue;
								if (exprMap.TryGetValueInCurrentScope(exprString, opValue))
								{
									exprMap.Add(GetExprString(instr.LeftOperand.Var), opValue);
									iter.Current->Delete();
									changed = true;
									continue;
								}
								else
								{
									exprMap.Add(exprString, instr.LeftOperand);
								}
							}
						}
						else
						{
							// replace operands
							for (auto & operand : instr.Operands)
							{
								Operand op;
								if (operand.IsVariable()&&exprMap.TryGetValue(GetExprString(operand.Var), op))
								{
									operand = op;
									changed = true;
								}
							}
							// try evaluate
							Operand opReplace;
							if (instr.Func && instr.Func->TryEvaluate(opReplace, instr.Operands))
							{
								exprMap.Add(GetExprString(instr.LeftOperand.Var), opReplace);
								changed = true;
								instr.Func = 0;
								instr.Operands.Clear();
								instr.Operands.Add(opReplace);
							}
							else if (instr.Operands.Count()==1 && !instr.Func)
							{
								if (instr.LeftOperand.IsVariable())
									exprMap.Add(GetExprString(instr.LeftOperand.Var), instr.Operands[0]);
								continue;
							}
							// try match with existing values
							String exprString = GetExprString(instr);
							Operand opValue;
							bool isConst = instr.Operands.Count()==1 && instr.Func == 0 &&
								!instr.Operands[0].IsVariable();
							if (!isConst && exprMap.TryGetValue(exprString, opValue) && !instr.IsVolatile)
							{
								exprMap.Add(GetExprString(instr.LeftOperand.Var), opValue);
								changed = true;
								program->VarDefs[instr.LeftOperand.Var->Id] = 0;
								iter.Current->Delete();
								continue;
							}
							// define a new value
							if (instr.LeftOperand.IsVariable())
								exprMap.Add(exprString, instr.LeftOperand);
							else if (instr.Func == Operation::Store && !instr.IsVolatile)
							{
								Instruction loadInstr(Operation::Load, instr.Operands[0], instr.Operands[2]);
								exprMap.Add(GetExprString(loadInstr), instr.Operands[1]);
							}
						}
					}
					// Modify phi function arguments of CFG successors
					for (auto childNode : node->Exits)
					{
						if (!childNode)
							break;
						for (auto & instr : childNode->Code) 
							if (instr.Func == Operation::Phi)
								for (auto & op : instr.Operands)
								{
									Operand opReplace;
									if (exprMap.TryGetValue(GetExprString(op.Var), opReplace)
										&& opReplace.IsVariable())
									{
										op = opReplace;
										changed = true;
									}
								}
					}
				}
				// recur in DomChildren
				bool nodeIsBranch = node->GetExitCount()==2 && node->Code.Count()>0 && node->Code.Last().Func == Operation::Branch && node->Code.Last().Operands[0].IsVariable();
				for (auto & childNode : node->DomChildren)
				{
					// propagate branch specific variables
					bool shouldPopScope = false;
					if (nodeIsBranch)
					{
						auto branchVar = node->Code.Last().Operands[0].Var;
						int branchType = node->Code.Last().Operands[2].IntValue;
						
						int value = 0;
						if (branchType > 1)
							throw InvalidProgramException(L"unknown branch type. This optimizer only recognizes JT and JF.");
						if (childNode == node->Exits[0])
						{
							if (branchType == 0) // jtrue
								value = 1;
							else
								value = 0;
							shouldPopScope = true;
							exprMap.PushScope();
							exprMap.Add(GetExprString(branchVar), value);
						}
						else if (childNode == node->Exits[1])
						{
							if (branchType == 0) // jtrue
								value = 0;
							else
								value = 1;
							shouldPopScope = true;
							exprMap.PushScope();
							exprMap.Add(GetExprString(branchVar), value);
						}
					}
					changed |= TraverseNode(program, childNode, exprMap);
					if (shouldPopScope)
						exprMap.PopScope();
				}
				exprMap.PopScope();
				return changed;
			}
		public:
			virtual ProgramOptimizationResult Optimize(RefPtr<ControlFlowGraph> program) override
			{
				ProgramOptimizationResult rs;
				ScopeDictionary<String, Operand> exprMap;
				rs.Changed = TraverseNode(program.Ptr(), program->Source, exprMap);
				rs.Program = program;
				return rs;
			}
		};
		IntraProcOptimizer * CreateUselessInstructionOptimizer()
		{
			return new UselessInstructionOptimizer();
		}
	}
}